package ace.module.oauth2.server.core.impl.authentication.password;

import java.security.Principal;
import java.util.Collections;
import java.util.HashMap;
import java.util.Map;
import java.util.Optional;
import org.apache.commons.lang3.StringUtils;
import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import org.springframework.boot.autoconfigure.condition.ConditionalOnBean;
import org.springframework.security.authentication.AuthenticationProvider;
import org.springframework.security.core.Authentication;
import org.springframework.security.core.AuthenticationException;
import org.springframework.security.core.userdetails.UserDetails;
import org.springframework.security.oauth2.core.*;
import org.springframework.security.oauth2.core.endpoint.OAuth2AuthorizationRequest;
import org.springframework.security.oauth2.core.oidc.OidcIdToken;
import org.springframework.security.oauth2.core.oidc.OidcScopes;
import org.springframework.security.oauth2.core.oidc.endpoint.OidcParameterNames;
import org.springframework.security.oauth2.jwt.Jwt;
import org.springframework.security.oauth2.server.authorization.OAuth2Authorization;
import org.springframework.security.oauth2.server.authorization.OAuth2AuthorizationService;
import org.springframework.security.oauth2.server.authorization.OAuth2TokenType;
import org.springframework.security.oauth2.server.authorization.authentication.OAuth2AccessTokenAuthenticationToken;
import org.springframework.security.oauth2.server.authorization.authentication.OAuth2ClientAuthenticationToken;
import org.springframework.security.oauth2.server.authorization.client.RegisteredClient;
import org.springframework.security.oauth2.server.authorization.context.AuthorizationServerContextHolder;
import org.springframework.security.oauth2.server.authorization.token.DefaultOAuth2TokenContext;
import org.springframework.security.oauth2.server.authorization.token.OAuth2TokenContext;
import org.springframework.security.oauth2.server.authorization.token.OAuth2TokenGenerator;
import org.springframework.stereotype.Component;
import org.springframework.util.Assert;

/**
 * @author caspar
 * @date 2023/2/8 13:43 密码模式 token
 */
@ConditionalOnBean(value = {Oauth2PasswordAuthenticationUserDetailsService.class})
@Component
public class Oauth2PasswordAuthenticationProvider implements AuthenticationProvider {
  private static final String ERROR_URI =
      "https://datatracker.ietf.org/doc/html/rfc6749#section-5.2";
  private static final OAuth2TokenType ID_TOKEN_TOKEN_TYPE =
      new OAuth2TokenType(OidcParameterNames.ID_TOKEN);
  private final Log log = LogFactory.getLog(getClass());
  private final OAuth2AuthorizationService authorizationService;
  private final OAuth2TokenGenerator<? extends OAuth2Token> tokenGenerator;
  private final Oauth2PasswordAuthenticationUserDetailsService userDetailsService;

  public Oauth2PasswordAuthenticationProvider(
      OAuth2AuthorizationService authorizationService,
      OAuth2TokenGenerator<? extends OAuth2Token> tokenGenerator,
      Oauth2PasswordAuthenticationUserDetailsService userDetailsService) {
    Assert.notNull(authorizationService, "authorizationService cannot be null");
    Assert.notNull(tokenGenerator, "tokenGenerator cannot be null");
    Assert.notNull(userDetailsService, "userDetailsService cannot be null");
    this.authorizationService = authorizationService;
    this.tokenGenerator = tokenGenerator;
    this.userDetailsService = userDetailsService;
  }

  @Override
  public Authentication authenticate(Authentication authentication) throws AuthenticationException {
    Oauth2PasswordAuthenticationToken oAuth2PasswordAuthenticationToken =
        (Oauth2PasswordAuthenticationToken) authentication;

    OAuth2ClientAuthenticationToken clientPrincipal =
        this.getAuthenticatedClientElseThrowInvalidClient(oAuth2PasswordAuthenticationToken);
    RegisteredClient registeredClient = clientPrincipal.getRegisteredClient();

    if (this.log.isTraceEnabled()) {
      this.log.trace("Retrieved registered client");
    }

    Optional<UserDetails> aceUserOptional =
        this.userDetailsService.loadUserDetails(oAuth2PasswordAuthenticationToken);

    if (StringUtils.equals(aceUserOptional.get().getPassword(), aceUserOptional.get().getPassword())
        == false) {
      throw new OAuth2AuthenticationException(
          new OAuth2Error(
              Oauth2PasswordAuthenticationErrorEnum.NOT_FOUND.getCode(),
              Oauth2PasswordAuthenticationErrorEnum.NOT_FOUND.getDesc(),
              ""));
    }

    UserDetails user = this.checkUser(aceUserOptional);

    OAuth2AuthorizationRequest authorizationRequest =
        OAuth2AuthorizationRequest.authorizationCode()
            .authorizationUri("http://" + registeredClient.getClientId())
            .clientId(registeredClient.getClientId())
            .redirectUri("")
            .scopes(registeredClient.getScopes())
            .state("")
            .additionalParameters(oAuth2PasswordAuthenticationToken.getAdditionalParameters())
            .build();

    OAuth2Authorization authorization =
        authorizationBuilder(
                registeredClient, oAuth2PasswordAuthenticationToken, authorizationRequest)
            .build();

    // @formatter:off
    DefaultOAuth2TokenContext.Builder tokenContextBuilder =
        DefaultOAuth2TokenContext.builder()
            .registeredClient(registeredClient)
            .principal(authorization.getAttribute(Principal.class.getName()))
            .authorizationServerContext(AuthorizationServerContextHolder.getContext())
            .authorization(authorization)
            .authorizedScopes(authorization.getAuthorizedScopes())
            .authorizationGrantType(AuthorizationGrantType.PASSWORD)
            .authorizationGrant(oAuth2PasswordAuthenticationToken);
    // @formatter:on

    OAuth2Authorization.Builder authorizationBuilder = OAuth2Authorization.from(authorization);

    // ----- Access token -----
    OAuth2TokenContext tokenContext =
        tokenContextBuilder.tokenType(OAuth2TokenType.ACCESS_TOKEN).build();
    OAuth2Token generatedAccessToken = this.tokenGenerator.generate(tokenContext);
    if (generatedAccessToken == null) {
      OAuth2Error error =
          new OAuth2Error(
              OAuth2ErrorCodes.SERVER_ERROR,
              "The token generator failed to generate the access token.",
              ERROR_URI);
      throw new OAuth2AuthenticationException(error);
    }

    if (this.log.isTraceEnabled()) {
      this.log.trace("Generated access token");
    }

    OAuth2AccessToken accessToken =
        new OAuth2AccessToken(
            OAuth2AccessToken.TokenType.BEARER,
            generatedAccessToken.getTokenValue(),
            generatedAccessToken.getIssuedAt(),
            generatedAccessToken.getExpiresAt(),
            tokenContext.getAuthorizedScopes());
    if (generatedAccessToken instanceof ClaimAccessor) {
      authorizationBuilder.token(
          accessToken,
          (metadata) ->
              metadata.put(
                  OAuth2Authorization.Token.CLAIMS_METADATA_NAME,
                  ((ClaimAccessor) generatedAccessToken).getClaims()));
    } else {
      authorizationBuilder.accessToken(accessToken);
    }

    // ----- Refresh token -----
    OAuth2RefreshToken refreshToken = null;
    if (registeredClient.getAuthorizationGrantTypes().contains(AuthorizationGrantType.REFRESH_TOKEN)
        &&
        // Do not issue refresh token to public client
        !clientPrincipal.getClientAuthenticationMethod().equals(ClientAuthenticationMethod.NONE)) {

      tokenContext = tokenContextBuilder.tokenType(OAuth2TokenType.REFRESH_TOKEN).build();
      OAuth2Token generatedRefreshToken = this.tokenGenerator.generate(tokenContext);
      if (!(generatedRefreshToken instanceof OAuth2RefreshToken)) {
        OAuth2Error error =
            new OAuth2Error(
                OAuth2ErrorCodes.SERVER_ERROR,
                "The token generator failed to generate the refresh token.",
                ERROR_URI);
        throw new OAuth2AuthenticationException(error);
      }

      if (this.log.isTraceEnabled()) {
        this.log.trace("Generated refresh token");
      }

      refreshToken = (OAuth2RefreshToken) generatedRefreshToken;
      authorizationBuilder.refreshToken(refreshToken);
    }

    // ----- ID token -----
    OidcIdToken idToken;
    if (authorizationRequest.getScopes().contains(OidcScopes.OPENID)) {
      // @formatter:off
      tokenContext =
          tokenContextBuilder
              .tokenType(ID_TOKEN_TOKEN_TYPE)
              .authorization(
                  authorizationBuilder
                      .build()) // ID token customizer may need access to the access token and/or
              // refresh token
              .build();
      // @formatter:on
      OAuth2Token generatedIdToken = this.tokenGenerator.generate(tokenContext);
      if (!(generatedIdToken instanceof Jwt)) {
        OAuth2Error error =
            new OAuth2Error(
                OAuth2ErrorCodes.SERVER_ERROR,
                "The token generator failed to generate the ID token.",
                ERROR_URI);
        throw new OAuth2AuthenticationException(error);
      }

      if (this.log.isTraceEnabled()) {
        this.log.trace("Generated id token");
      }

      idToken =
          new OidcIdToken(
              generatedIdToken.getTokenValue(),
              generatedIdToken.getIssuedAt(),
              generatedIdToken.getExpiresAt(),
              ((Jwt) generatedIdToken).getClaims());
      authorizationBuilder.token(
          idToken,
          (metadata) ->
              metadata.put(OAuth2Authorization.Token.CLAIMS_METADATA_NAME, idToken.getClaims()));
    } else {
      idToken = null;
    }

    authorization = authorizationBuilder.build();

    this.authorizationService.save(authorization);

    if (this.log.isTraceEnabled()) {
      this.log.trace("Saved authorization");
    }

    Map<String, Object> additionalParameters = Collections.emptyMap();
    if (idToken != null) {
      additionalParameters = new HashMap<>();
      additionalParameters.put(OidcParameterNames.ID_TOKEN, idToken.getTokenValue());
    }

    if (this.log.isTraceEnabled()) {
      this.log.trace("Authenticated token input");
    }

    return new OAuth2AccessTokenAuthenticationToken(
        registeredClient, clientPrincipal, accessToken, refreshToken, additionalParameters);
  }

  private UserDetails checkUser(Optional<UserDetails> aceUserOptional) {
    if (aceUserOptional.isEmpty()) {
      throw new OAuth2AuthenticationException(
          new OAuth2Error(
              Oauth2PasswordAuthenticationErrorEnum.NOT_FOUND.getCode(),
              Oauth2PasswordAuthenticationErrorEnum.NOT_FOUND.getDesc(),
              ""));
    }

    if (aceUserOptional.get().isEnabled() == false) {
      throw new OAuth2AuthenticationException(
          new OAuth2Error(
              Oauth2PasswordAuthenticationErrorEnum.ACCOUNT_DISABLED.getCode(),
              Oauth2PasswordAuthenticationErrorEnum.ACCOUNT_DISABLED.getDesc(),
              ""));
    }

    if (aceUserOptional.get().isAccountNonExpired() == false) {
      throw new OAuth2AuthenticationException(
          new OAuth2Error(
              Oauth2PasswordAuthenticationErrorEnum.ACCOUNT_EXPIRED.getCode(),
              Oauth2PasswordAuthenticationErrorEnum.ACCOUNT_EXPIRED.getDesc(),
              ""));
    }

    if (aceUserOptional.get().isAccountNonLocked() == false) {
      throw new OAuth2AuthenticationException(
          new OAuth2Error(
              Oauth2PasswordAuthenticationErrorEnum.ACCOUNT_LOCKED.getCode(),
              Oauth2PasswordAuthenticationErrorEnum.ACCOUNT_LOCKED.getDesc(),
              ""));
    }

    if (aceUserOptional.get().isCredentialsNonExpired() == false) {
      throw new OAuth2AuthenticationException(
          new OAuth2Error(
              Oauth2PasswordAuthenticationErrorEnum.CREDENTIALS_EXPIRED.getCode(),
              Oauth2PasswordAuthenticationErrorEnum.CREDENTIALS_EXPIRED.getDesc(),
              ""));
    }
    return aceUserOptional.get();
  }

  @Override
  public boolean supports(Class<?> authentication) {
    return Oauth2PasswordAuthenticationToken.class.isAssignableFrom(authentication);
  }

  private OAuth2ClientAuthenticationToken getAuthenticatedClientElseThrowInvalidClient(
      Authentication authentication) {
    OAuth2ClientAuthenticationToken clientPrincipal = null;
    if (OAuth2ClientAuthenticationToken.class.isAssignableFrom(
        authentication.getPrincipal().getClass())) {
      clientPrincipal = (OAuth2ClientAuthenticationToken) authentication.getPrincipal();
    }
    if (clientPrincipal != null && clientPrincipal.isAuthenticated()) {
      return clientPrincipal;
    }
    throw new OAuth2AuthenticationException(OAuth2ErrorCodes.INVALID_CLIENT);
  }

  private OAuth2Authorization.Builder authorizationBuilder(
      RegisteredClient registeredClient,
      Authentication principal,
      OAuth2AuthorizationRequest authorizationRequest) {
    return OAuth2Authorization.withRegisteredClient(registeredClient)
        .principalName(principal.getName())
        .authorizationGrantType(AuthorizationGrantType.PASSWORD)
        .attribute(Principal.class.getName(), principal)
        .attribute(OAuth2AuthorizationRequest.class.getName(), authorizationRequest);
  }
}
